'''
Created on Jun 17, 2009

@author: user
'''
import os
import unittest
import cliques.mdp
import cProfile
import pstats
import graph
from cliques.cliqueFinder import cliqueFinder
from cliques.mdp.mdpBasedSearch import mdpBasedSearch
from cliques.mdp.heuristicMdp import heuristicMdp
import cliques.experiment_runner
import cliques.utils
import tests.mdpAbstractTest
class unittest_mdp(tests.mdpAbstractTest.mdpAbstractTest):

    def testMdpStateActions(self):
        ''' Check that MDP state returns the correct possible actions '''
        num_nodes = 6
        k = 5
        ckg = graph.graph()
        ckg.generate(num_nodes, num_nodes*(num_nodes-1)/2)      
        generated=[]
        state = cliques.mdp.mdpSolver.cliqueSearchState(ckg, generated)
        assert(len(list(state.actions()))==0)
        
        nodes =ckg.nodes()
        generated=[nodes[0]]
        state = cliques.mdp.mdpSolver.cliqueSearchState(ckg, generated)
        assert len(list(state.actions()))==len(generated)
        assert nodes[0] in state.actions()
    
    def testIsGoal(self):
        num_nodes = 3
        k = 3
        ckg = graph.graph()
        ckg.generate(num_nodes, num_nodes*(num_nodes-1)/2)
        old_nodes = list(ckg.nodes())
        new_nodes = ['n0','n1']  
        ckg.add_nodes(new_nodes) 
        ckg.add_edge(old_nodes[0],new_nodes[0])
        ckg.add_edge(old_nodes[0],new_nodes[1])                   
        generated=new_nodes
        generated.append(old_nodes[1])
        state = cliques.mdp.mdpSolver.cliqueSearchState(ckg, generated)
        solver = cliques.mdp.mdpSolver.cliqueMdpSolver(3, 0.7, num_nodes+len(new_nodes))
        assert solver.is_goal(state)

    def testBug1(self):
        num_nodes=5
        k=3
        nodes=['n0','n1','n2','n3','n4']
        ckg = graph.graph()
        ckg.add_nodes(nodes)
        ckg.add_edge(nodes[1],nodes[2])
        ckg.add_edge(nodes[2],nodes[3])
        ckg.add_edge(nodes[2],nodes[0])
        ckg.add_edge(nodes[3],nodes[4])
        generated = [nodes[4]]
        solver = cliques.mdp.mdpSolver.cliqueMdpSolver(3,0.4,num_nodes)
        (v, policy) = solver.run(ckg,generated)
        
        assert v.has_key(state)
        assert policy.has_key(state)

    def testSmartMdp(self):
        mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.cliqueMdpSolver,"MDP")
        smart_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smartCliqueMdpSolver,"smart-MDP")
        algorithms = [mdp_based_search,smart_mdp_based_search]
        assertion_funcs = [self.assert_halt,self.assert_current_node, self.assert_less_states]
        
        # Profile
        self.clique_search_run_comparator(5, 3, 7, algorithms, assertion_funcs)
        #cProfile.run("clique_search_run_comparator(self,5, 3, 7, algorithms, assertion_funcs)")
      
   

    def testBugGraphX(self):
        ''' Bug  '''
        BUG_FILE_NAME = "BUG-GRAPH[3,0]-6.7.graph"        
        k=3
        start_node=u'0'
        smart_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smartCliqueMdpSolver,"smart-MDP")
        smarter_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smarterCliqueMdpSolver,"smarter-MDP")
        algorithms = [smart_mdp_based_search, smarter_mdp_based_search]
        assertion_funcs = [self.assert_halt,self.assert_current_node, self.assert_less_states]
        
        graph_file = open(BUG_FILE_NAME,"r+")
        read_xml = graph_file.read()
        graph_file.close()
        test_graph = graph.graph()
        test_graph.read(read_xml)    
        self.simulate_run(test_graph, algorithms, assertion_funcs, k, start_node)    
        
    def testSmartMdpChooseNode(self):
        ''' Wierd choice of SmartMdp '''
        # Prepare the solver        
        smart_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smartCliqueMdpSolver,"smart-MDP")
        
        BUG_FILE_NAME = "BUG-GRAPH[3,0]-6.7.graph"        
        k=3
        start_node=u'0'
        graph_file = open(BUG_FILE_NAME,"r+")
        read_xml = graph_file.read()
        graph_file.close()
        test_graph = graph.graph()
        test_graph.read(read_xml)    
        smart_mdp_based_search.setup(test_graph, k, start_node)
        # Get to BUG state
        smart_mdp_based_search.step()
        
        # Chose node
        chosen_node = smart_mdp_based_search.choose_node()
        assert chosen_node!=u'3'


    def testSmarterMdpWrong(self):
        graph_file = open("SmarterWrong - BugGraph.graph","r+")
        read_xml = graph_file.read()
        graph_file.close()
        test_graph = graph.graph()
        test_graph.read(read_xml)

        start_node = u'1'
        k = 4
        smart_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smartCliqueMdpSolver,"smart-MDP")
        smarter_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smarterCliqueMdpSolver,"smarter-MDP")

        algorithms = [smart_mdp_based_search,smarter_mdp_based_search]        
        assertion_funcs = [self.assert_halt,self.assert_current_node, self.assert_less_states]
        
        self.compare(test_graph, k, k+1, algorithms, assertion_funcs)

    def testSmarterMdp(self):
        smart_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smartCliqueMdpSolver,"smart-MDP")
        smarter_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smarterCliqueMdpSolver,"smarter-MDP")
        algorithms = [smart_mdp_based_search, smarter_mdp_based_search]
        assertion_funcs = [self.assert_halt,self.assert_current_node, self.assert_less_states]
        max_iterations = 25
        for nodes in xrange(5,7):
            for iteration in xrange(max_iterations):
                print " ---- Nodes %d, Iteration%d/%d" % (nodes,iteration,max_iterations)
                self.clique_search_run_comparator(nodes, 3, 7, algorithms, assertion_funcs)

    def testHeuristicMdp(self):   
        ''' Verify that heuristic MDP works like smarter MDP (since both are optimal) '''     
        smarter_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smarterCliqueMdpSolver,"smarter-MDP")
        heuristic_mdp_based_search = mdpBasedSearch(heuristicMdp,"heuristic-MDP")        
        algorithms = [smarter_mdp_based_search, heuristic_mdp_based_search]
        assertion_funcs = [self.assert_halt,self.assert_current_node, self.assert_less_states]
        max_iterations = 25
        for nodes in xrange(5,7):
            for iteration in xrange(max_iterations):
                print " ---- Nodes %d, Iteration %d/%d" % (nodes,iteration,max_iterations)
                self.clique_search_run_comparator(nodes, 3, 7, algorithms, assertion_funcs)
        

    def testUnlimitedLookahead(self):
        ''' Verify that unlimited lookahead is equal to heuristic MDP '''
        smarter_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smarterCliqueMdpSolver,"smarter-MDP")
        heuristic_mdp_based_search = mdpBasedSearch(heuristicMdp,"heuristic-MDP")        
        algorithms = [smarter_mdp_based_search, heuristic_mdp_based_search]
        assertion_funcs = [self.assert_halt,self.assert_current_node, self.assert_less_states]
        max_iterations = 25
        for nodes in xrange(5,7):
            for iteration in xrange(max_iterations):
                print " ---- Nodes %d, Iteration %d/%d" % (nodes,iteration,max_iterations)
                self.clique_search_run_comparator(nodes, 3, 7, algorithms, assertion_funcs)
                

    def testHeuristicBug(self):
        GRAPH_FILE = "BUG-GRAPH[4,1]-unittest-heuristicMdp.graph"
        bug_graph = cliques.utils.import_graph(GRAPH_FILE)
        start_node = u'1'
        k=4
        smarter_mdp_based_search = mdpBasedSearch(cliques.mdp.mdpSolver.smarterCliqueMdpSolver,"smarter-MDP")
        heuristic_mdp_based_search = mdpBasedSearch(heuristicMdp,"heuristic-MDP")        
        algorithms = [smarter_mdp_based_search, heuristic_mdp_based_search]
        assertion_funcs = [self.assert_halt,self.assert_current_node, self.assert_less_states]
        
        self.simulate_run(bug_graph, algorithms, assertion_funcs, k, start_node)
        

    def testLookaheadLevels(self):
        ''' Verify that unlimited lookahead is equal to heuristic MDP '''


    def testMdpNextStates(self):
        # Check full graph that have been completely generated
        num_nodes = 5
        k = 5
        ckg = graph.graph()
        ckg.generate(num_nodes, num_nodes*(num_nodes-1)/2)
        
        nodes =ckg.nodes()
        last_node = nodes[0]        
        generated=[last_node]
        state = cliques.mdp.mdpSolver.cliqueSearchState(ckg, generated)
        solver = cliques.mdp.mdpSolver.cliqueMdpSolver(k, 1, num_nodes)
        
        next_states = list(solver.generate_next_states(state, last_node))
        assert len(next_states)==1
        
        # Check an almost full graph with 2 nodes unexplored
        state.generated=[nodes[0],nodes[1]]        
        state.ckg.del_edge(nodes[0],nodes[1]) # Remove the edge between them since is unknown yet
        solver.p = 0.2
        next_states = list(solver.generate_next_states(state,nodes[0]))

        # Check that one state has an edge between them and another state without
        if next_states[0][0].ckg.has_edge(nodes[0],nodes[1]):
            assert next_states[0][1]==solver.p
            assert next_states[1][1]==1-solver.p
            assert next_states[1][0].ckg.has_edge(nodes[0],nodes[1])==False
        else:
            assert next_states[0][1]==1-solver.p
            assert next_states[1][1]==solver.p 
            assert next_states[1][0].ckg.has_edge(nodes[0],nodes[1])
       
       
    def testMdpNextStates2(self):
        num_nodes = 4
        k = 3
        p=0.3
        ckg = graph.graph()
        nodes = ['n0','n1','n2','n3']
        ckg.add_node(nodes[0])
        ckg.add_node(nodes[1])
        ckg.add_node(nodes[2])
        ckg.add_node(nodes[3])
        ckg.add_edge(nodes[0],nodes[1])
        ckg.add_edge(nodes[0],nodes[2])
        ckg.add_edge(nodes[0],nodes[3])
        generated = [nodes[1],nodes[2],nodes[3]]
        
        state = cliques.mdp.mdpSolver.cliqueSearchState(ckg, generated)        
        solver = cliques.mdp.mdpSolver.cliqueMdpSolver(k, p, num_nodes)
                
        #
        
        action = nodes[1]
        
        state_to_prob = dict()
        next_states = list(solver.generate_next_states(state, action))
        
        new_generated = list([nodes[2],nodes[3]])        
        new_ckg = graph.graph()
        new_ckg.add_graph(ckg)
        new_ckg.add_edge(nodes[1],nodes[2])
        new_ckg.add_edge(nodes[1],nodes[3])        
        new_state = cliques.mdp.mdpSolver.cliqueSearchState(new_ckg,new_generated)
        state_to_prob[new_state] = pow(p,2)
               
        new_ckg = graph.graph()
        new_ckg.add_graph(ckg)
        new_ckg.add_edge(nodes[1],nodes[2])       
        new_state = cliques.mdp.mdpSolver.cliqueSearchState(new_ckg,new_generated)
        state_to_prob[new_state] = p*(1-p)

        new_ckg = graph.graph()
        new_ckg.add_graph(ckg)
        new_ckg.add_edge(nodes[1],nodes[3])       
        new_state = cliques.mdp.mdpSolver.cliqueSearchState(new_ckg,new_generated)
        state_to_prob[new_state] = p*(1-p)

        new_ckg = graph.graph()
        new_ckg.add_graph(ckg)      
        new_state = cliques.mdp.mdpSolver.cliqueSearchState(new_ckg,new_generated)
        state_to_prob[new_state] = pow((1-p),2)
        
        for state,prob in next_states:
            same_states = [same_state for same_state in state_to_prob.keys() if same_state.same(state)]
            assert(len(same_states)>0)
            same_state = same_states[0]
            assert state_to_prob.has_key(same_state)
            assert state_to_prob[same_state]==prob
                


if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testMdpNextStates']
#    cProfile.run("unittest.main()","mdp-test2.profile")    
#    profile_stats = pstats.Stats("mdp-test2.profile")
#    profile_stats.strip_dirs()
#    profile_stats.sort_stats('time')
#    profile_stats.print_stats(15)
    #import sys;sys.argv = ['', 'unittest_mdp.testHeuristicBug']
    unittest.main()
    
    
    